import torch
import torch.nn as nn

class FNN(nn.Module):
    def __init__(self, layers, activation, dropout_rate=0.001, in_tf=None, out_tf=None):
        """
        初始化 FNN 模型，添加 Dropout 层以评估模型参数不确定性。

        参数：
        layers (list): 每层神经元数量的列表，例如 [4, 64, 64, 1]
        activation (callable): 激活函数，例如 torch.nn.ReLU()
        dropout_rate (float): Dropout 概率，默认值为 0.1
        in_tf (callable, optional): 输入变换函数
        out_tf (callable, optional): 输出变换函数
        """
        super().__init__()
        self.activation = activation
        self.linears = nn.ModuleList()
        self.dropouts = nn.ModuleList()  # 添加 Dropout 层
        self.in_tf = in_tf
        self.out_tf = out_tf
        self.dropout_rate = dropout_rate

        # 构建网络层
        for i in range(1, len(layers)):
            self.linears.append(nn.Linear(layers[i-1], layers[i]))
            # 在每个线性层后（除了最后一层）添加 Dropout
            if i < len(layers) - 1:
                self.dropouts.append(nn.Dropout(p=dropout_rate))
            else:
                self.dropouts.append(nn.Identity())  # 最后一层不加 Dropout

            # 权重初始化
            nn.init.xavier_uniform_(self.linears[-1].weight)
            nn.init.zeros_(self.linears[-1].bias)

    def forward(self, inputs, mc_dropout=False):
        """
        前向传播，支持蒙特卡洛 Dropout 进行模型参数不确定性评估。

        参数：
        inputs (torch.Tensor): 输入张量
        mc_dropout (bool): 是否启用蒙特卡洛 Dropout，True 时在推理阶段也启用 Dropout

        返回：
        torch.Tensor: 模型输出
        """
        X = inputs

        # 输入变换
        if self.in_tf:
            X = self.in_tf(X)

        # 确保 Dropout 在推理阶段也能生效（如果 mc_dropout=True）
        if mc_dropout:
            self.train()  # 启用训练模式，使 Dropout 生效
        else:
            self.eval()  # 正常推理模式

        # 线性层和 Dropout
        for linear, dropout in zip(self.linears[:-1], self.dropouts[:-1]):
            X = linear(X)
            X = self.activation(X)
            X = dropout(X)  # 应用 Dropout

        # 最后一层（无激活函数和 Dropout）
        X = self.linears[-1](X)

        # 输出变换
        if self.out_tf:
            X = self.out_tf(X)

        return X
